{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import Tensor\n",
    "\n",
    "from typing import List, Optional, Union"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "___\n",
    "**input**: $\\gamma$ (lr), $\\theta_{0}$ (params), $f(\\theta)$ (objective), $\\lambda $ (weight decay), $\\mu$ (momentum), $\\tau$ (dampening), $nesterov$, $maximize$\n",
    "___\n",
    "$for \\quad t = 1 \\quad to ... do \\\\$\n",
    "$\\qquad g_{t} \\gets \\nabla_{\\theta}f_{t}(\\theta_{t-1}) \\\\$\n",
    "$\\qquad if \\quad \\lambda \\neq 0: \\\\$\n",
    "$\\qquad \\qquad g_{t} \\gets g_{t}+\\lambda \\theta_{t-1} \\\\$  \n",
    "$\\qquad if \\quad \\mu \\neq 0: \\\\$\n",
    "$\\qquad \\qquad if \\quad t > 1: \\\\$\n",
    "$\\qquad \\qquad \\qquad b_{t} \\gets \\mu b_{t-1}+(1-\\tau)g_{t} \\\\$\n",
    "$\\qquad \\qquad else \\\\$\n",
    "$\\qquad \\qquad \\qquad b_{t}\\gets g_{t} \\\\$\n",
    "$\\qquad \\qquad if \\quad nesterov \\\\$\n",
    "$\\qquad \\qquad \\qquad g_{t}\\gets g_{t}+\\mu b_{t} \\\\$\n",
    "$\\qquad \\qquad else \\\\$\n",
    "$\\qquad \\qquad \\qquad g_{t}\\gets b_{t} \\\\$\n",
    "$\\qquad if \\quad maximize: \\\\$\n",
    "$\\qquad \\qquad \\theta_{t} \\gets \\theta_{t-1} + \\gamma g_{t} \\\\$\n",
    "$\\qquad else: \\\\$\n",
    "$\\qquad \\qquad \\theta_{t} \\gets \\theta_{t-1}-\\gamma g_{t} \\\\$ \n",
    "___\n",
    "return $\\theta_{t}$\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _single_tensor_sgd(\n",
    "    params: List[Tensor],\n",
    "    grads: List[Tensor],\n",
    "    momentum_buffer_list: List[Optional[Tensor]],\n",
    "    grad_scale: Optional[Tensor],\n",
    "    found_inf: Optional[Tensor],\n",
    "    weight_decay: float,\n",
    "    momentum: float,\n",
    "    lr: float,\n",
    "    dampening: float,\n",
    "    nesterov: bool,\n",
    "    maximize: bool,\n",
    "):\n",
    "    assert grad_scale is None and found_inf is None\n",
    "\n",
    "    for i, param in enumerate(params):\n",
    "        grad = grads[i] if not maximize else -grads[i]\n",
    "\n",
    "        if weight_decay != 0:\n",
    "            grad = grad.add(param, alpha=weight_decay)\n",
    "\n",
    "        if momentum != 0:\n",
    "            buf = momentum_buffer_list[i]\n",
    "\n",
    "            if buf is None:\n",
    "                buf = torch.clone(grad).detach()\n",
    "                momentum_buffer_list[i] = buf\n",
    "            else:\n",
    "                buf.mul_(momentum).add_(grad, alpha=1 - dampening)\n",
    "\n",
    "            if nesterov:\n",
    "                grad = grad.add(buf, alpha=momentum)\n",
    "            else:\n",
    "                grad = buf\n",
    "\n",
    "        param.add_(grad, alpha=-lr)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "___\n",
    "**input**: $\\gamma$ (lr), $\\beta_{1}$, $\\beta_{2}$ (betas), $\\theta_{0}$ (params), $f(\\theta)$ (objective), $\\lambda $ (weight decay), $amsgrad$, $maximize$, $\\epsilon$ (epsilon)\n",
    "\n",
    "**initialize**: $m_{0} \\gets 0$ (first moment), $v_{0} \\gets 0$ (second moment), $\\hat{v_{0}}^{max} \\gets 0$\n",
    "___\n",
    "$for \\quad t=1 \\quad to ... do \\\\$\n",
    "$\\qquad if \\quad maximize: \\\\$\n",
    "$\\qquad \\qquad g_{t} \\gets - \\nabla_{\\theta}f_{t}(\\theta_{t-1}) \\\\$\n",
    "$\\qquad else: \\\\$\n",
    "$\\qquad \\qquad g_{t} \\gets \\nabla_{\\theta}f_{t}(\\theta_{t-1}) \\\\$\n",
    "$\\qquad if \\quad \\lambda \\neq 0: \\\\$\n",
    "$\\qquad \\qquad g_{t} = g_{t} + \\lambda \\theta_{t-1} \\\\$\n",
    "$\\qquad m_{t} \\gets \\beta_{1} m_{t-1} + (1-\\beta_{1})g_{t} \\\\$\n",
    "$\\qquad v_{t} \\gets \\beta_{2} v_{t-1} + (1-\\beta_{2})g^{2}_{t} \\\\$\n",
    "$\\qquad \\hat{m_{t}} \\gets m_{t}/(1-\\beta^{t}_{1}) \\\\$\n",
    "$\\qquad \\hat{v_{t}} \\gets v_{t}/(1-\\beta^{t}_{2}) \\\\$\n",
    "$\\qquad if \\quad amsgrad \\\\$\n",
    "$\\qquad \\qquad \\hat{v_{t}}^{max} \\gets max(\\hat{v_{t-1}}^{max}, \\hat{v_{t}}) \\\\$\n",
    "$\\qquad \\qquad \\theta_{t} \\gets \\theta_{t-1} - \\gamma\\hat{m_{t}} / (\\sqrt{\\hat{v_{t}}^{max}} + \\epsilon) \\\\$\n",
    "$\\qquad else \\\\$\n",
    "$\\qquad \\qquad \\theta_{t} \\gets \\theta_{t-1} - \\gamma\\hat{m_{t}} / (\\sqrt{\\hat{v_{t}}} + \\epsilon) \\\\$\n",
    "___\n",
    "return $\\theta_{t}$\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 删除了 amsgrad 操作，amsgrad 操作需要变换到复数域\n",
    "def _single_tensor_adam(\n",
    "    params: List[Tensor],\n",
    "    grads: List[Tensor],\n",
    "    exp_avgs: List[Tensor],\n",
    "    exp_avg_sqs: List[Tensor],\n",
    "    state_steps: List[Tensor],\n",
    "    beta1: Union[float, Tensor],\n",
    "    beta2: Union[float, Tensor],\n",
    "    lr: Union[float, Tensor],\n",
    "    weight_decay: float,\n",
    "    eps: float,\n",
    "):\n",
    "    for i, param in enumerate(params):\n",
    "        grad = grads[i] \n",
    "        exp_avg = exp_avgs[i]\n",
    "        exp_avg_sq = exp_avg_sqs[i]\n",
    "        step_t = state_steps[i]\n",
    "\n",
    "        # update step\n",
    "        step_t += 1\n",
    "\n",
    "        if weight_decay != 0:\n",
    "            grad = grad.add(param, alpha=weight_decay)\n",
    "\n",
    "        device_beta1 = beta1\n",
    "\n",
    "        # Decay the first and second moment running average coefficient\n",
    "        exp_avg.lerp_(grad, 1 - device_beta1)\n",
    "\n",
    "        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)\n",
    "\n",
    "        step = step_t\n",
    "\n",
    "        bias_correction1 = 1 - beta1**step\n",
    "        bias_correction2 = 1 - beta2**step\n",
    "\n",
    "        step_size = lr / bias_correction1\n",
    "\n",
    "        bias_correction2_sqrt = bias_correction2**0.5\n",
    "\n",
    "        denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)\n",
    "\n",
    "        param.addcdiv_(exp_avg, denom, value=-step_size)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "___\n",
    "**input**: $\\gamma$ (lr), $\\beta_{1}$, $\\beta_{2}$ (betas), $\\theta_{0}$ (params), $f(\\theta)$ (objective), $\\lambda $ (weight decay), $amsgrad$, $maximize$, $\\epsilon$ (epsilon)\n",
    "\n",
    "**initialize**: $m_{0} \\gets 0$ (first moment), $v_{0} \\gets 0$ (second moment), $\\hat{v_{0}}^{max} \\gets 0$\n",
    "___\n",
    "$for \\quad t=1 \\quad to ... do \\\\$\n",
    "$\\qquad if \\quad maximize: \\\\$\n",
    "$\\qquad \\qquad g_{t} \\gets - \\nabla_{\\theta}f_{t}(\\theta_{t-1}) \\\\$\n",
    "$\\qquad else: \\\\$\n",
    "$\\qquad \\qquad g_{t} \\gets \\nabla_{\\theta}f_{t}(\\theta_{t-1}) \\\\$\n",
    "$\\qquad \\theta_{t} \\gets \\theta_{t-1} - \\gamma \\lambda \\theta_{t-1} \\\\$\n",
    "$\\qquad m_{t} \\gets \\beta_{1} m_{t-1} + (1-\\beta_{1})g_{t} \\\\$\n",
    "$\\qquad v_{t} \\gets \\beta_{2} v_{t-1} + (1-\\beta_{2})g^{2}_{t} \\\\$\n",
    "$\\qquad \\hat{m_{t}} \\gets m_{t}/(1-\\beta^{t}_{1}) \\\\$\n",
    "$\\qquad \\hat{v_{t}} \\gets v_{t}/(1-\\beta^{t}_{2}) \\\\$\n",
    "$\\qquad if \\quad amsgrad \\\\$\n",
    "$\\qquad \\qquad \\hat{v_{t}}^{max} \\gets max(\\hat{v_{t-1}}^{max}, \\hat{v_{t}}) \\\\$\n",
    "$\\qquad \\qquad \\theta_{t} \\gets \\theta_{t-1} - \\gamma\\hat{m_{t}} / (\\sqrt{\\hat{v_{t}}^{max}} + \\epsilon) \\\\$\n",
    "$\\qquad else \\\\$\n",
    "$\\qquad \\qquad \\theta_{t} \\gets \\theta_{t-1} - \\gamma\\hat{m_{t}} / (\\sqrt{\\hat{v_{t}}} + \\epsilon) \\\\$\n",
    "___\n",
    "return $\\theta_{t}$\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _single_tensor_adamw(\n",
    "    params: List[Tensor],\n",
    "    grads: List[Tensor],\n",
    "    exp_avgs: List[Tensor],\n",
    "    exp_avg_sqs: List[Tensor],\n",
    "    state_steps: List[Tensor],\n",
    "    beta1: Union[Tensor, float],\n",
    "    beta2: Union[Tensor, float],\n",
    "    lr: Union[Tensor, float],\n",
    "    weight_decay: float,\n",
    "    eps: float,\n",
    "    maximize: bool,\n",
    "):\n",
    "    for i, param in enumerate(params):\n",
    "        grad = grads[i] if not maximize else -grads[i]\n",
    "        exp_avg = exp_avgs[i]\n",
    "        exp_avg_sq = exp_avg_sqs[i]\n",
    "        step_t = state_steps[i]\n",
    "\n",
    "        # update step\n",
    "        step_t += 1\n",
    "\n",
    "        # Perform stepweight decay\n",
    "        param.mul_(1 - lr * weight_decay)\n",
    "        \n",
    "        device_beta1 = beta1\n",
    "\n",
    "        # Decay the first and second moment running average coefficient\n",
    "        exp_avg.lerp_(grad, 1 - device_beta1)\n",
    "        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)\n",
    "\n",
    "        step = step_t\n",
    "\n",
    "        bias_correction1 = 1 - beta1**step\n",
    "        bias_correction2 = 1 - beta2**step\n",
    "\n",
    "        step_size = lr / bias_correction1\n",
    "\n",
    "        bias_correction2_sqrt = bias_correction2**0.5\n",
    "\n",
    "        denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)\n",
    "\n",
    "        param.addcdiv_(exp_avg, denom, value=-step_size)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
